#!/usr/bin/env python3
# yolo_ocr_hsv_optimized_node.py

# ==============================================================================
# Section 1: Imports and Path Resolution
# ==============================================================================
import os
import sys
import time
import json
import functools
import warnings
from collections import Counter, deque
from typing import Deque, Optional, Tuple, List, Union, Dict, Any

import rospy
from std_msgs.msg import String

import cv2
import numpy as np
import torch
from ultralytics import YOLO

# Try to import CnOcr, but make it optional so the node can run without it.
try:
    from cnocr import CnOcr
except ImportError:
    CnOcr = None
    rospy.logwarn("cnocr library not found. OCR functionality will be disabled.")

warnings.filterwarnings("ignore")

def _resolve_paths() -> Tuple[str, str, str]:
    """
    Resolves key directory paths based on the script's location.
    Assumes a standard project structure where this script is in a 'scripts' folder.
    """
    try:
        # Assumes: .../nav_ws/src/pkg/scripts/ -> .../nav_ws/
        repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
    except Exception:
        repo_root = os.getcwd() # Fallback to current directory
        rospy.logwarn(f"Could not resolve repo root from file path. Using fallback: {repo_root}")

    detect_dir = os.path.join(repo_root, 'yolo-jetson', 'detect')
    rs2_dir = os.path.join(repo_root, 'yolo-jetson', 'rs2')
    return repo_root, detect_dir, rs2_dir

REPO_ROOT, DETECT_DIR, RS2_DIR = _resolve_paths()
if DETECT_DIR not in sys.path:
    sys.path.append(DETECT_DIR)
if RS2_DIR not in sys.path:
    sys.path.append(RS2_DIR)

# GStreamerWrapper is a dependency from the original node
try:
    from GStreamerWrapper import GStreamerWrapper
except ImportError:
    rospy.logerr("Fatal: Failed to import GStreamerWrapper. Make sure yolo-jetson/rs2 is in your PYTHONPATH.")
    sys.exit(1)

# ==============================================================================
# Section 2: Optimized Pipeline Logic (from your provided code)
# ==============================================================================
class OptimizedYoloOcrHsvPipeline:
    """Optimized YOLO + OCR + HSV Pipeline"""
    COLOR_CLASSES = {"red", "orange", "yellow", "green", "blue"}
    NUMBER_CLASSES = {"one", "two", "three", "four", "five", "six"}
    STATE_CLASSES = {"down", "normal", "up"}

    COLOR_RANGES = {
        "red": [np.array([[0, 20, 225], [4, 188, 255]], dtype=np.uint8)],
        "orange": [
            np.array([[0, 97, 222], [11, 155, 255]], dtype=np.uint8),
            np.array([[18, 97, 222], [179, 155, 255]], dtype=np.uint8)
        ],
        "yellow": [np.array([[23, 122, 219], [32, 193, 255]], dtype=np.uint8)],
        "green": [np.array([[37, 46, 207], [47, 185, 255]], dtype=np.uint8)],
        "blue": [np.array([[92, 171, 222], [108, 239, 255]], dtype=np.uint8)],
    }
    DIGIT_TO_NAME = {"1": "one", "2": "two", "3": "three", "4": "four", "5": "five", "6": "six"}

    def __init__(self, model_path: str, device: Optional[str] = None, debug: bool = False, enable_warmup: bool = True):
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"Model weights not found: {model_path}")

        self.debug = debug
        self.device = device if device is not None else ('cuda' if torch.cuda.is_available() else 'cpu')
        
        self.model = YOLO(model_path)
        if hasattr(self.model, 'model'):
            self.model.model.eval()
            
        if enable_warmup:
            self._warmup_model()
        
        self._init_cnocr()
        
        self._morph_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        self._last_crop_offsets = {"top": 0, "left": 0}
        
        if self.debug:
            rospy.loginfo(f"Optimized pipeline initialized - Device: {self.device}, CnOCR: {'Available' if self.ocr_available else 'Unavailable'}")

    def _warmup_model(self):
        try:
            dummy = np.full((640, 640, 3), 128, dtype=np.uint8)
            with torch.no_grad():
                _ = self.model.predict(dummy, verbose=False, conf=0.1, device=self.device, save=False, show=False)
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if self.debug:
                rospy.loginfo("Pipeline internal model warm-up complete.")
        except Exception as e:
            rospy.logwarn(f"Pipeline internal model warm-up failed: {e}")

    def _init_cnocr(self):
        self.ocr_available = False
        self.cnocr = None
        if CnOcr is not None:
            try:
                self.cnocr = CnOcr(
                    model_name='number-densenet_lite_136-fc',
                    det_model_name='en_PP-OCRv3_det',
                    cand_alphabet=['1', '2', '3', '4', '5', '6']
                )
                self.ocr_available = True
            except Exception as e:
                rospy.logerr(f"Failed to initialize CnOcr: {e}")
        else:
            self.ocr_available = False

    def process_image(self, image: np.ndarray, top_cut_ratio: float = 0.25, width_center_ratio: float = 0.5, first_stage_conf: float = 0.6, final_conf: float = 0.6) -> Dict[str, Any]:
        """Main image processing function for the optimized pipeline."""
        # ... (The full implementation of your OptimizedYoloOcrHsvPipeline methods)
        # ... (crop_image, yolo_first_stage, analyze_first_stage_results, etc.)
        # This code is omitted for brevity but should be pasted here from your provided script.
        # For this example to be runnable, the full class methods must be included.
        # A placeholder implementation is provided below.

        # --- Placeholder Implementation ---
        # Replace this with the actual full methods from your optimized script.
        try:
            cropped = self.crop_image(image, top_cut_ratio, width_center_ratio)
            first_stage_results = self.yolo_first_stage(cropped, conf=first_stage_conf)
            detections = self.analyze_first_stage_results(first_stage_results, conf=first_stage_conf)
            roi_list = self.build_adaptive_rois(detections, cropped)
            ocr_result = self.recognize_number_with_ocr_optimized(roi_list)
            hsv_result = self.hsv_cone_color_vote(roi_list)
            final_results = self.finalize_results(first_stage_results, detections, ocr_result, hsv_result, conf=final_conf)
            
            cone_color, number, dashboard_state = "None", "None", "None"
            for class_name, detected_name, _ in final_results:
                if class_name == "cone": cone_color = detected_name
                elif class_name == "number": number = detected_name
                elif class_name == "dashboard": dashboard_state = detected_name

            return {
                'cone_color': cone_color, 'number': number, 'dashboard_state': dashboard_state,
                'processing_info': {'status': 'success'} if self.debug else {}
            }
        except Exception as e:
            return {
                'cone_color': "None", 'number': "None", 'dashboard_state': "None",
                'error': str(e), 'processing_info': {}
            }

    # You MUST paste all the other helper methods from your `OptimizedYoloOcrHsvPipeline` class here.
    # e.g., crop_image, yolo_first_stage, analyze_first_stage_results, build_adaptive_rois,
    # _find_best_detection, _make_color_masks_optimized, _predict_color_by_hsv_optimized,
    # hsv_cone_color_vote, _preprocess_for_ocr_optimized, recognize_number_with_ocr_optimized,
    # _extract_ocr_result, _extract_digit_candidate_bbox_optimized, finalize_results, _find_best_from_results

# ==============================================================================
# Section 3: The Adapted ROS Node
# ==============================================================================

def majority_vote(cache: Deque[Tuple[str, str, str]]) -> Optional[Tuple[str, str, str]]:
    """Uses majority voting to determine the final result from the cache."""
    if not cache:
        return None
    merged = [f'{t[0]}-{t[1]}-{t[2]}' for t in cache]
    top, _ = Counter(merged).most_common(1)[0]
    c, n, d = top.split('-')
    return c, n, d

class YOLOOcrHsvOptimizedNode:
    """
    ROS node for the optimized YOLO+OCR+HSV pipeline.
    Inherits the ROS interface from the original node but uses the new optimized processing pipeline.
    """
    def __init__(self) -> None:
        rospy.init_node('yolo_ocr_hsv_optimized_node', anonymous=False)

        # --- Parameters (same as original node) ---
        self.width = rospy.get_param('~width', 1920)
        self.height = rospy.get_param('~height', 1080)
        self.first_stage_conf = rospy.get_param('~first_stage_conf', 0.6)
        self.final_conf = rospy.get_param('~final_conf', 0.6)
        self.top_cut_ratio = rospy.get_param('~top_cut_ratio', 0.25)
        self.width_center_ratio = rospy.get_param('~width_center_ratio', 0.5)
        self.cache_size = int(rospy.get_param('~cache_size', 10))
        self.target_fps = int(rospy.get_param('~target_fps', 12))
        self.debug_mode = rospy.get_param('~debug', False)

        # --- Model Path Discovery (same as original node) ---
        candidate_models = [
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0811', 'best.pt'),
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0808', 'best.pt'),
            os.path.join(DETECT_DIR, 'weights', 'all_yolo_all_epochs_0805', 'best.pt'),
        ]
        self.model_path = rospy.get_param('~model_path', '')
        if not self.model_path:
            for p in candidate_models:
                if os.path.exists(p):
                    self.model_path = p
                    rospy.loginfo(f'Found model: {p}')
                    break
        if not self.model_path or not os.path.exists(self.model_path):
            rospy.logerr('No model weights found. Set ~model_path param to a valid .pt file.')
            raise RuntimeError('model_path not found')

        # --- Component Initialization (MODIFIED) ---
        rospy.loginfo("Initializing optimized YOLO pipeline...")
        self.pipeline = OptimizedYoloOcrHsvPipeline(
            model_path=self.model_path,
            debug=self.debug_mode,
            enable_warmup=False  # Node handles warmup with real frames
        )
        rospy.loginfo("Optimized YOLO pipeline initialized.")

        self.gs = GStreamerWrapper(width=self.width, height=self.height)

        self._warmup_engine()

        # --- ROS Publishers/Subscribers (same as original node) ---
        self.result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self.cmd_sub = rospy.Subscriber('/yolo_detect', String, self._cmd_cb, queue_size=10)
        self.detailed_result_pub = rospy.Publisher('/yolo_detailed_result', String, queue_size=10)

        # --- State Management (same as original node) ---
        self._running = False
        self._cache: Deque[Tuple[str, str, str]] = deque(maxlen=self.cache_size)
        self._frame_interval = 1.0 / max(1, self.target_fps)
        self._last_time = 0.0
        self._frame_count = 0
        self._success_count = 0

        self._timer = rospy.Timer(rospy.Duration(0.01), self._on_timer)

        rospy.loginfo('Optimized YOLO node ready. Send "yolostart" on /yolo_detect to begin.')
        rospy.loginfo(f'Model: {self.model_path}')
        rospy.loginfo(f'Resolution: {self.width}x{self.height}')
        rospy.loginfo(f'Target FPS: {self.target_fps}')
        rospy.loginfo(f'OCR Available: {self.pipeline.ocr_available}')

    def _warmup_engine(self) -> None:
        """Warms up the YOLO engine using a real or dummy frame."""
        rospy.loginfo("Starting YOLO engine warm-up...")
        start_time = time.time()
        frame = None
        while time.time() - start_time < 5.0: # Increased timeout for camera
            frame = self.gs.GetFrame()
            if frame is not None:
                rospy.loginfo("Acquired real frame for warm-up.")
                break
            time.sleep(0.1)

        if frame is None:
            frame = np.zeros((self.height, self.width, 3), dtype=np.uint8)
            rospy.logwarn('Using dummy frame for warm-up as camera failed to provide one.')
        
        try:
            _ = self.pipeline.process_image(
                frame,
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                final_conf=self.final_conf,
            )
            rospy.loginfo('YOLO engine warm-up completed.')
        except Exception as e:
            rospy.logerr(f'YOLO engine warm-up failed: {e}')
            raise

    def _cmd_cb(self, msg: String) -> None:
        """Handles command messages."""
        cmd = (msg.data or '').strip().lower()
        if cmd == 'yolostart' and not self._running:
            self._cache.clear()
            self._running = True
            self._frame_count = 0
            self._success_count = 0
            rospy.loginfo('YOLO OCR+HSV Optimized started')
        elif cmd in ('stop', 'yolostop') and self._running:
            self._running = False
            self._cache.clear()
            rospy.loginfo('YOLO OCR+HSV Optimized stopped')
        elif cmd == 'status':
            rospy.loginfo(f'Status: running={self._running}, frames={self._frame_count}, success={self._success_count}')

    def _on_timer(self, _evt) -> None:
        """Timer callback function for processing image frames."""
        if not self._running or rospy.is_shutdown():
            return

        frame = self.gs.GetFrame()
        if frame is None:
            return

        now = time.time()
        if now - self._last_time < self._frame_interval:
            return
        self._last_time = now
        self._frame_count += 1

        try:
            res = self.pipeline.process_image(
                frame,
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                final_conf=self.final_conf,
            )

            cone = res.get('cone_color', 'None')
            num = res.get('number', 'None')
            dash = res.get('dashboard_state', 'None')

            if cone != 'None' and num != 'None' and dash != 'None':
                self._cache.append((cone, num, dash))
                self._success_count += 1

            # Publish detailed result as a JSON string
            detailed_msg = {
                'frame_count': self._frame_count,
                'success_count': self._success_count,
                'cache_size': len(self._cache),
                'cone_color': cone,
                'number': num,
                'dashboard_state': dash,
                'processing_info': res.get('processing_info', {}),
                'timestamp': now
            }
            self.detailed_result_pub.publish(String(data=json.dumps(detailed_msg)))

            if len(self._cache) >= self.cache_size:
                final_triplet = majority_vote(self._cache)
                if final_triplet:
                    fc, fn, fd = final_triplet
                    msg = f'{fc}-{fn}-{fd}'
                    self.result_pub.publish(String(data=msg))
                    rospy.loginfo(f'Published /yolo_result: {msg}')
                    rospy.loginfo(f'Processing complete: {self._frame_count} frames, {self._success_count} successful detections.')
                
                self._running = False
                self._cache.clear()

        except Exception as e:
            rospy.logerr(f'Error processing frame {self._frame_count}: {e}')

if __name__ == '__main__':
    try:
        node = YOLOOcrHsvOptimizedNode()
        rospy.spin()
    except Exception as e:
        rospy.logerr(f'Optimized YOLO node failed to initialize: {e}')
        import traceback
        traceback.print_exc()